
'''
1 str -> json
2 json -> values
3 compare values VS search_values
4 return satisfy, sim
'''
from flask import Flask, request, jsonify, current_app

def create_app():
    import os
    from util_for_clipvitb32 import clip_vit_b32
    app = Flask(__name__)
    
    # 在应用创建时立即加载模型（每个进程启动时执行）
    with app.app_context():
        app.clip_model = clip_vit_b32()
        print(f"Model loaded at: {id(app.clip_model)}, PID: {os.getpid()}")
    return app

# 创建全局应用实例（Gunicorn 需要）
app = create_app()

# 初始化并全局绑定 clip_model，用于后续 Flask 并发请求
# @app.before_first_request
# def load_model():
#     from util_for_clipvitb32 import clip_vit_b32
#     current_app.clip_model = clip_vit_b32()
#     print(f"Model loaded at: {id(current_app.clip_model)}")

def get_satisfy_sim(target_json_dict_str, non_empty_keys,
                    search_joined_values_list, search_embs, 
                    threshold, device):
    import json,torch
    import torch.nn.functional as F
    from demo_qwen3_api import get_matched_info_by_jsondict
    from util_for_clipvitb32 import get_embedding
    try:
        target_json_dict = json.loads(target_json_dict_str)
    except Exception:
        target_json_dict = {}

    joined_values_list = get_matched_info_by_jsondict(target_json_dict, non_empty_keys)

    sim_list = []
    satisfy = True

    # 确保 search_embs 中每个元素都是 torch.Tensor
    # 并且全部移动到正确的 device
    device_obj = torch.device(device)
    search_embs = [
        torch.tensor(emb).to(device_obj) if not isinstance(emb, torch.Tensor) else emb.to(device_obj)
        for emb in search_embs
    ]

    # iterate over pairs (joined_value, search_joined_value, search_emb_tensor)
    for i, (joined_value, search_joined_value) in \
                        enumerate(zip(joined_values_list, search_joined_values_list)):

        search_emb_tensor = search_embs[i]  # already on device

        # 数值匹配逻辑
        if search_joined_value.isdigit() or \
            (search_joined_value.startswith('-') and search_joined_value[1:].isdigit()):
            
            search_num = int(search_joined_value)

            if search_num == 0:
                if joined_value == "":
                    sim_list.append(1.0)
                    continue
                if joined_value.isdigit() or (joined_value.startswith('-') and joined_value[1:].isdigit()):
                    joined_num = int(joined_value)
                    if joined_num in [0, -1]:
                        sim_list.append(1.0)
                        continue
                    else:
                        satisfy = False
                        break
                else:
                    satisfy = False
                    break
            if joined_value.isdigit() or \
                (joined_value.startswith('-') and joined_value[1:].isdigit()):

                joined_num = int(joined_value)
                sim = 1.0 if search_num == joined_num else 0.0
                if sim < threshold:
                    satisfy = False
                    break
                sim_list.append(sim)
                continue
            else:
                satisfy = False
                break

        # 特殊键 visual/view_types 的集合包含逻辑
        cur_keys = non_empty_keys[joined_values_list.index(joined_value)] if joined_value in joined_values_list else []
        if len(cur_keys) >= 2 and cur_keys[-2:] == ["visual", "view_types"]:
            sv_set = set([v.strip() for v in search_joined_value.split(",")])
            jv_set = set([v.strip() for v in joined_value.split(",")])
            if not sv_set.issubset(jv_set):
                satisfy = False
                break
            sim_list.append(1.0)
            continue

        # 否则用 CLIP 文本相似度（GPU）
        joined_values_emb = get_embedding(current_app.clip_model, joined_value)

        # convert to torch tensor on device if needed
        if isinstance(joined_values_emb, torch.Tensor):
            joined_tensor = joined_values_emb.to(device)
        else:
            # numpy or list -> tensor
            joined_tensor = torch.tensor(joined_values_emb, 
                                        dtype=torch.float32, device=device)

        a = search_emb_tensor.unsqueeze(0)
        b = joined_tensor.unsqueeze(0)
        search_sim = F.cosine_similarity(a, b).item()

        if search_sim < threshold:
            satisfy = False
            break
        sim_list.append(search_sim)

    return satisfy, sim_list

@app.route('/satisfy_sim', methods=['POST'])
def satisfy_sim():
    import torch
    # 解析输入参数
    data = request.json
    if not data:
        return jsonify({'error': 'No input data'}), 400

    target_json_dict_str = data.get("target_json_dict_str", "")
    non_empty_keys = data.get("non_empty_keys", [])
    search_joined_values_list = data.get("search_joined_values_list", [])
    search_embs = data.get("search_embs", [])
    threshold = float(data.get("threshold", 0.3))
    device = data.get("device", "cpu")


    satisfy, sim_list = get_satisfy_sim(
        target_json_dict_str,
        non_empty_keys,
        search_joined_values_list,
        search_embs,
        threshold,
        device
    )
    return jsonify({"satisfy": satisfy, "sim_list": sim_list})

if __name__ == '__main__':    
    import os
    os.environ['GRADIO_TEMP_DIR'] = '/mnt/nas/shengjie/tmp'
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('-c', '--cuda', type=str, default='0', help='CUDA device index or device string, eg "0" or "cuda:0"')
    parser.add_argument('-p', '--port', type=int, default=20023, help='gradio port')
    args, unknown = parser.parse_known_args()

    os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda

    app.run(host="0.0.0.0", port=args.port)
    
# 使用以下命令用 4 个 worker 开启 gunicorn 运行 Flask 应用
# (假定flask_retrival.py中的app对象名为app)
# 在终端中运行:
# gunicorn -w 4 -b 0.0.0.0:20022 flask_retrival:app

